#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <complex.h>
#include <mpi.h>

#define PI 3.14159265358979323846

typedef double complex cplx;

// 创建MPI数据类型用于复数打包
void create_complex_type(MPI_Datatype *cplx_type) {
    int blocklengths[2] = {1, 1};
    MPI_Datatype types[2] = {MPI_DOUBLE, MPI_DOUBLE};
    MPI_Aint offsets[2];
    
    cplx dummy;
    MPI_Get_address(&dummy, &offsets[0]);
    MPI_Get_address(&imag(dummy), &offsets[1]);
    offsets[1] -= offsets[0];
    offsets[0] = 0;
    
    MPI_Type_create_struct(2, blocklengths, offsets, types, cplx_type);
    MPI_Type_commit(cplx_type);
}

// ... (保留之前的fft函数和main函数结构)

int main(int argc, char **argv) {
    MPI_Init(&argc, &argv);
    
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    
    // 创建自定义复数类型
    MPI_Datatype cplx_type;
    create_complex_type(&cplx_type);
    
    // ... (其余初始化代码与之前相同)
    
    // 使用打包数据进行通信
    double pack_start = MPI_Wtime();
    
    // 计算打包缓冲区大小
    int position = 0;
    int pack_size;
    MPI_Pack_size(local_size, cplx_type, MPI_COMM_WORLD, &pack_size);
    char *pack_buffer = malloc(pack_size);
    
    // 在发送方打包数据
    if (rank == 0) {
        for (int i = 0; i < local_size; i++) {
            MPI_Pack(&global_data[i], 1, cplx_type, 
                    pack_buffer, pack_size, &position, MPI_COMM_WORLD);
        }
    }
    
    // 分发打包数据
    MPI_Scatter(pack_buffer, pack_size, MPI_PACKED,
               pack_buffer, pack_size, MPI_PACKED,
               0, MPI_COMM_WORLD);
    
    // 接收方解包数据
    position = 0;
    cplx *local_data = malloc(local_size * sizeof(cplx));
    for (int i = 0; i < local_size; i++) {
        MPI_Unpack(pack_buffer, pack_size, &position,
                  &local_data[i], 1, cplx_type, MPI_COMM_WORLD);
    }
    
    double pack_end = MPI_Wtime();
    
    if (rank == 0) {
        printf("Data packing time: %.6f seconds\n", pack_end - pack_start);
    }
    
    // ... (其余计算和收集代码与之前相同)
    
    MPI_Type_free(&cplx_type);
    free(pack_buffer);
    MPI_Finalize();
    return 0;
}